"""
词向量和数字向量的转换
"""


class WordSequence:
    UNK_TAG = "UNK"
    PAD_TAG = "PAD"
    EOS_TAG = "EOS"
    SOS_TAG = "SOS"

    UNK = 0
    PAD = 1
    EOS = 2
    SOS = 3

    dict = {
        UNK_TAG: UNK,
        PAD_TAG: PAD,
        EOS_TAG: EOS,
        SOS_TAG: SOS
    }

    count = {}

    reverse_dict = {}

    def fit(self, sentence):
        """
        统计词频
        :param sentence: 词语列表
        :return:
        """
        for word in sentence:
            self.count[word] = self.count.get(word, 0) + 1

    def build_vocab(self, min_count=5, max_count=None, max_feature=None):
        """
        创建字典
        :param min_count: 最小词频
        :param max_count: 最大词频
        :param max_feature: 最大词语数量
        :return:
        """
        temp = self.count.copy()  # 深拷贝，直接赋值时浅拷贝
        if min_count:
            for key in temp:
                if self.count.get(key, 0) < min_count:
                    del self.count[key]

        if max_count:
            for key in temp:
                if self.count.get(key, max_count+1) > max_count:
                    del self.count[key]

        if max_feature:
            self.count =dict(sorted(self.count.items(), key=lambda x: x[1],
                                    reverse=True)[: max_feature])

        for key in self.count.keys():
            self.dict[key] = len(self.dict)

        self.reverse_dict = dict(zip(self.dict.values(), self.dict.keys()))

    def transform(self, sentence, max_length=None, add_eos=False):
        """
        词向量转数字向量
        :param sentence: 词向量
        :param max_length: 最大长度
        :param add_eos: 是否添加结束符
        :return:
        """
        if len(sentence) > max_length:
            sentence = sentence[: max_length]

        sentence_len = len(sentence)

        if add_eos:
            sentence.append(self.EOS_TAG)

        if sentence_len < max_length:
            sentence += [self.PAD_TAG] * (max_length - sentence_len)

        return [self.dict.get(i, self.UNK) for i in sentence]

    def reverse_transform(self, sentence):
        """
        将数字向量转换为词向量
        :param sentence: 数字向量
        :return:
        """
        result = []
        for num in sentence:
            if num == self.EOS:
                break
            result.append(self.reverse_dict.get(num, self.UNK_TAG))

        return result

    def __len__(self):
        return len(self.dict)






